from sklearn.datasets import load_iris
import pandas as pd
from sklearn.metrics import accuracy_score, classification_report
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.tree import plot_tree
import matplotlib.pyplot as plt


# 决策树
def desition_tree_iris():
    """
    用决策树对鸢尾花进行分类
    :return:
    """
    # 1. 获取数据集
    iris = load_iris()
    # 2. 划分数据集
    x_train, x_test, y_train, y_test = train_test_split(iris.data, iris.target, random_state=22)
    # 3. 特征工程
    # 4. 模型训练(决策树预估器）
    estimator = DecisionTreeClassifier(criterion="entropy")
    estimator.fit(x_train, y_train)
    # 5. 模型评估
    y_predict = estimator.predict(x_test)
    print("预测结果为：", y_predict)
    print(f"直接对比真实值和预测值：\n{y_test == y_predict}")
    print("测试集的准确率:", estimator.score(x_test, y_test))

    # 可视化
    # 可以直接用Matplotlib绘制决策树
    plt.figure(figsize=(10, 10))
    plot_tree(estimator, filled=True)
    plt.savefig("../images/iris_tree.png")
    plt.show()


def titanic_demo():
    """
    泰坦尼克号生存预测案例
    步骤：
    1. 数据加载与初步观察
    2. 数据预处理（缺失值处理、特征编码）
    3. 划分训练集和测试集
    4. 决策树模型训练
    5. 模型评估与结果分析

    :return: None
    """
    # 1. 数据加载与初步观察
    try:
        data = pd.read_csv("../data/titanic.csv")

        # 输出数据基本信息
        print("=" * 50)
        print("数据集基本信息：")
        print(f"数据集形状：{data.shape}")
        print("\n前5行数据：")
        print(data.head())
        print("\n缺失值统计：")
        print(data.isnull().sum())

    except FileNotFoundError:
        print("错误：未找到数据集文件，请检查路径！")
        return

    # 2. 数据预处理
    print("\n" + "=" * 50)
    print("数据预处理阶段：")

    # 2.1 处理缺失值
    data['age'].fillna(data['age'].median(), inplace=True)
    print("\n已处理年龄缺失值，填充中位数：", data['age'].median())

    # 2.2 特征编码
    # 性别编码（male:0, female:1）
    data['sex'] = LabelEncoder().fit_transform(data['sex'])

    # 舱位等级编码（'1st':0, '2nd':1, '3rd':2）
    data['pclass'] = data['pclass'].map({'1st': 0, '2nd': 1, '3rd': 2})
    print("\n舱位等级编码结果：1st→0, 2nd→1, 3rd→2")

    # 3. 划分训练集和测试集
    features = ["pclass", "age", "sex"]
    target = "survived"

    x_train, x_test, y_train, y_test = train_test_split(
        data[features],
        data[target],
        test_size=0.2,
        random_state=22
    )

    # 4. 模型训练
    estimator = DecisionTreeClassifier(
        criterion="entropy",
        max_depth=3,
        random_state=22
    )
    estimator.fit(x_train, y_train)

    # 5. 模型评估
    y_predict = estimator.predict(x_test)

    print("\n测试集准确率：", accuracy_score(y_test, y_predict))
    print("\n分类报告：")
    print(classification_report(y_test, y_predict, target_names=['未生存', '生存']))

    # 可视化决策树
    plt.figure(figsize=(12, 8))
    plot_tree(estimator,
              feature_names=features,
              class_names=['未生存', '生存'],
              filled=True)
    plt.savefig("../images/titanic_tree.png")
    plt.show()


if __name__ == '__main__':
    # desition_tree_iris()
    titanic_demo()
